"""
内容：输入-预测数值模型示例
日期：2020年7月6日
作者：Howie
"""
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
import matplotlib.pyplot as plt

# 预设
weight = 2.0
bias = 3.0

# 超参
EPOCH = 50
BATCH_SIZE = 64


def dataset_generating(verbose=False):
    """
    # 生成数据集
    :return:
    """
    # 线性回归模型：Y = 2X + 0.25
    X_train = np.random.random((1000, 1))
    Y_train = X_train * weight + np.random.random((1000, 1)) / bias

    X_test = np.random.random((100, 1))
    Y_test = X_test * weight + np.random.random((100, 1)) / bias

    if verbose:
        print(
            "Train Samples: {}, Test Samples: {}".format(
                X_train.shape,
                X_test.shape))
        plt.figure()
        plt.title('Samples Visualization')
        plt.plot(X_train, Y_train, 'ro')
        plt.plot(X_test, Y_test, 'bo')
        plt.legend(['Train', 'Test'], loc='upper left')
        plt.savefig('./logs/Samples.pdf')
        plt.show()

    return X_train, X_test, Y_train, Y_test


def model_building(mlp=True, deep=True):
    """
    # 模型准备
    :return: model
    """
    model = Sequential()
    if mlp:
        model.add(Dense(units=64, input_dim=1, activation='relu'))
        model.add(Dense(units=1))
    if mlp and deep:
        model.add(Dense(units=64, input_dim=1, activation='relu'))
        model.add(Dense(units=64, activation='relu'))
        model.add(Dense(units=1))
    else:
        model.add(Dense(units=1, input_dim=1))
    return model


def model_training():
    """
    # 训练模型
    :return:
    """
    # 生成数据集
    X_train, X_test, Y_train, Y_test = dataset_generating()
    # 搭建模型
    models = [model_building(mlp=True, deep=False),
              model_building(mlp=True, deep=True),
              model_building(mlp=False, deep=False)]
    titles = ['Multi-Layer Perceptron',
              'Deep Multi-Layer Perceptron',
              'Perceptron']
    for title_idx, model in enumerate(models):
        # 设置模型训练过程
        model.compile(optimizer='rmsprop', loss='mse')
        # 训练模型
        hist = model.fit(X_train, Y_train, epochs=EPOCH, batch_size=BATCH_SIZE)
        # 查看训练过程
        plt.plot(hist.history['loss'])
        plt.ylim(0.0, 1.5)
        plt.ylabel('loss')
        plt.xlabel('epoch')
        plt.legend(['train'], loc='upper left')
        plt.title(titles[title_idx])
        plt.grid()
        plt.savefig('./logs/' + titles[title_idx] + '.pdf')
        plt.show()
        # 评价模型
        loss = model.evaluate(X_test, Y_test, batch_size=BATCH_SIZE)
        print("Loss: {}".format(loss))


if __name__ == '__main__':
    model_training()
